'''
基于Redis的分布式对话历史管理

Redis作为消息存储的优势：
1. 高性能：内存数据库，读写速度快
2. 持久化：支持数据持久化到磁盘
3. 分布式：支持多实例共享对话历史
4. 过期策略：自动清理过期数据
5. 数据结构丰富：支持多种数据格式

应用场景：
- 多用户对话系统
- 分布式部署的聊天服务
- 需要持久化对话历史的场景
- 跨会话保持对话连续性
'''
import os

# 先安装redis客户端
# pip install --upgrade --quiet redis

from dotenv import load_dotenv
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, ConfigurableFieldSpec
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的环境变量
load_dotenv()

# ========== 构建对话处理链 ==========
'''
创建聊天提示词模板，定义对话的结构：
- system: 设定AI的角色和能力
- MessagesPlaceholder: 历史消息占位符，运行时从Redis加载
- human: 当前用户输入
'''
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个优秀的助手，你的能力是{ability}，每次你返回问题的结果不要超过30个字"),
        MessagesPlaceholder(variable_name="history"),  # 历史消息将从Redis注入
        ("human", "{input}")  # 当前用户输入
    ]
)

# 初始化大语言模型
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen3-32B", streaming=True, temperature=0)

# 输出解析器：将模型输出转换为字符串
parser = StrOutputParser()

# 构建完整的处理链：提示词 -> 大模型 -> 输出解析
runnable = prompt | llm | parser

# ========== Redis历史消息管理 ==========
def get_message_history(user_id: str, conversation_id: str) -> RedisChatMessageHistory:
    """
    获取或创建Redis中的对话历史

    设计理念：多租户隔离
    - 每个用户有独立的对话空间
    - 每个用户可以有多轮对话（conversation_id）
    - 支持用户在不同设备间同步对话历史

    Redis键设计：user_id:conversation_id
    示例：
      "user_123:conversation_456" -> 用户123的对话456的历史
      "user_123:conversation_789" -> 用户123的对话789的历史

    参数：
      user_id: 用户唯一标识（如用户ID、用户名等）
      conversation_id: 对话会话标识（如客服会话ID、聊天窗口ID等）

    返回：
      RedisChatMessageHistory实例，用于操作该对话的历史消息
    """
    # 构建Redis存储键：使用冒号分隔的层级结构
    key = f"{user_id}:{conversation_id}"

    # 创建Redis聊天消息历史实例
    return RedisChatMessageHistory(
        session_id=key,        # Redis中的键名
        url="redis://localhost:6379/0",  # Redis连接URL
        ttl=3600               # 生存时间（秒），1小时后自动过期
    )

# ========== 配置带历史管理的执行器 ==========
'''
RunnableWithMessageHistory 的核心功能：
1. 自动管理对话历史
2. 根据配置获取对应的历史记录
3. 将历史注入到提示词中
4. 自动保存新的对话记录
'''
with_message_history = RunnableWithMessageHistory(
    # 基础处理链
    runnable,

    # 历史消息工厂函数：根据配置参数获取对应的历史记录
    get_message_history,

    # 配置参数映射
    input_messages_key="input",      # 用户输入在输入字典中的key
    history_messages_key="history",  # 历史消息在提示词模板中的key

    # 历史工厂的配置规范
    history_factory_config=[
        ConfigurableFieldSpec(
            id="user_id",           # 配置参数ID
            annotation=str,         # 参数类型
            name="User ID",         # 参数名称
            description="用户的唯一标识符。",  # 参数描述
            default="",             # 默认值
            is_shared=True,         # 是否在多个调用间共享
        ),
        ConfigurableFieldSpec(
            id="conversation_id",
            annotation=str,
            name="Conversation ID",
            description="对话的唯一标识符。",
            default="",
            is_shared=True,
        ),
    ],
)

# ========== 使用示例和测试代码 ==========
def demonstrate_redis_storage():
    """
    演示Redis存储功能的使用
    """
    print("🚀 Redis对话历史存储演示")
    print("=" * 50)

    # 测试对话1：用户A的客服对话
    print("\n💬 测试对话1：用户A的客服对话")
    try:
        response1 = with_message_history.invoke(
            {
                "ability": "客服咨询",
                "input": "你好，我想查询订单状态"
            },
            config={
                "configurable": {
                    "user_id": "user_001",
                    "conversation_id": "customer_service_001"
                }
            }
        )
        print(f"🤖 AI回复: {response1}")
    except Exception as e:
        print(f"❌ 错误: {e} (请确保Redis服务正在运行)")

    # 测试对话2：同一用户的不同对话
    print("\n💬 测试对话2：用户A的学习助手对话")
    try:
        response2 = with_message_history.invoke(
            {
                "ability": "学习辅导",
                "input": "请解释什么是机器学习"
            },
            config={
                "configurable": {
                    "user_id": "user_001",
                    "conversation_id": "study_assistant_001"
                }
            }
        )
        print(f"🤖 AI回复: {response2}")
    except Exception as e:
        print(f"❌ 错误: {e}")

    # 测试对话3：不同用户的对话（完全隔离）
    print("\n💬 测试对话3：用户B的对话")
    try:
        response3 = with_message_history.invoke(
            {
                "ability": "通用助手",
                "input": "今天天气怎么样"
            },
            config={
                "configurable": {
                    "user_id": "user_002",
                    "conversation_id": "general_chat_001"
                }
            }
        )
        print(f"🤖 AI回复: {response3}")
    except Exception as e:
        print(f"❌ 错误: {e}")

def check_redis_contents():
    """
    检查Redis中的存储内容（需要redis-cli或redis-py）
    """
    print("\n🔍 Redis存储结构检查")
    print("=" * 50)
    print("可以使用以下命令查看Redis内容:")
    print("1. redis-cli")
    print("2. KEYS '*'                    # 查看所有键")
    print("3. LRANGE 'user_001:customer_service_001' 0 -1  # 查看具体对话历史")
    print("4. TTL 'user_001:customer_service_001'     # 查看剩余过期时间")

# ========== Redis配置和优化建议 ==========
'''
📝 Redis生产环境配置建议：

1. 连接池配置：
   import redis
   pool = redis.ConnectionPool(host='localhost', port=6379, db=0, max_connections=10)
   RedisChatMessageHistory(..., connection_pool=pool)

2. 密码认证：
   url="redis://:password@localhost:6379/0"

3. 集群支持：
   url="redis://host1:port1,host2:port2,host3:port3/0"

4. TTL策略：
   - 短期会话：ttl=3600 (1小时)
   - 长期记忆：ttl=2592000 (30天)  
   - 永久存储：ttl=None

5. 键前缀：
   key = f"chat:user:{user_id}:conversation:{conversation_id}"
'''

class AdvancedRedisMessageManager:
    """
    高级Redis消息管理器示例
    """

    def __init__(self, redis_url="redis://localhost:6379/0"):
        self.redis_url = redis_url

    def get_message_history_with_prefix(self, user_id: str, conversation_id: str, prefix="chat") -> RedisChatMessageHistory:
        """
        使用前缀的Redis键管理
        """
        key = f"{prefix}:user:{user_id}:conversation:{conversation_id}"
        return RedisChatMessageHistory(
            session_id=key,
            url=self.redis_url,
            ttl=3600
        )

    def cleanup_user_conversations(self, user_id: str):
        """
        清理用户的所有对话（高级功能）
        """
        import redis
        r = redis.Redis.from_url(self.redis_url)

        # 使用模式匹配删除用户的所有对话
        pattern = f"*:{user_id}:*"
        keys = r.keys(pattern)
        if keys:
            r.delete(*keys)
            print(f"✅ 已清理用户 {user_id} 的 {len(keys)} 个对话")

    def get_conversation_list(self, user_id: str):
        """
        获取用户的所有对话列表（高级功能）
        """
        import redis
        r = redis.Redis.from_url(self.redis_url)

        pattern = f"*:{user_id}:*"
        conversations = r.keys(pattern)
        return [key.decode('utf-8') for key in conversations]

# 运行演示
if __name__ == "__main__":
    # 演示基本功能
    demonstrate_redis_storage()

    # 显示检查方法
    check_redis_contents()

    print("\n🎯 Redis存储的核心优势总结:")
    print("1. 💾 持久化存储：重启服务不丢失对话历史")
    print("2. 🔄 多实例共享：支持分布式部署")
    print("3. 👥 多用户隔离：用户数据完全独立")
    print("4. ⚡ 高性能：内存级读写速度")
    print("5. 🗑️ 自动清理：TTL自动过期机制")
    print("6. 🔍 可查询：支持历史记录检索和分析")